//------------------------------------------------------------------------------
// <copyright file="WinRtStringSocket.h" company="Microsoft">
//     Copyright (c) Microsoft Corporation.  All rights reserved.
// </copyright>
//------------------------------------------------------------------------------

#pragma once

#ifndef USE_RAPIDJSON

#include <ppltasks.h>
#include <functional>
#include <atomic>
#include "IRpcStringSocket.h"
#include "ppl.h"

using namespace Windows::Foundation;
using namespace Windows::Storage::Streams;
using namespace Windows::Storage::Streams;
using namespace Windows::Networking::Sockets;
using namespace Windows::Foundation;
using namespace concurrency;

namespace
{
	using namespace Microsoft::Internal::GamesTest::Rpc::Server;
	const IRpcStringSocket::DataReceivedCallback EmptyDataReceivedCallback = [](IRpcStringSocket*, std::wstring) { };
	const IRpcStringSocket::SocketDisconnectedCallback EmptySocketDisconnectedCallback = [](IRpcStringSocket*) { };
}

namespace Microsoft { namespace Internal { namespace GamesTest { namespace Rpc { namespace Server
{
	//
	// This class represents a client that is connected to the server.
	// For mocking the system's WinRT classes we use a technique called "Test Base Class Injection".
	// http://portals/engineering/Wiki/Wiki%20Pages/How%20To%20Mock%20An%20Api%20C++%202.aspx
	//
	template<typename Base>
	class WinRtStringSocketT : public Base, public IRpcStringSocket
	{
	private:
		DataWriter^ m_dataWriter;
		DataReader^ m_dataReader;
		StreamSocket^ m_socket;
		concurrency::cancellation_token_source m_cancelTokenSource;
		concurrency::critical_section m_sendDataLock;
		bool m_isListening;
		DataReceivedCallback m_dataReceivedCallback;
		SocketDisconnectedCallback m_clientDisconnectedCallback;

	public:
		virtual ~WinRtStringSocketT()
		{
			Shutdown();

			delete m_socket;
			delete m_dataReader;
			delete m_dataWriter;
		}

	public:
		//
		// Gets the IpAddress of the client that this socket connects to.
		//
		virtual std::wstring GetIpAddress()
		{
			// GetIpAddress is actually used only for debug spew when errors occur, so no need to optimize strings here.
			return std::wstring(m_socket->Information->RemoteAddress->DisplayName->Data(), m_socket->Information->RemoteAddress->DisplayName->Length());
		}

		//
		// Initializes a new instance of the WinRtStringSocket class.
		//	socket: The actual socket to be used for sending/receiving data.
		//
		WinRtStringSocketT(StreamSocket^ socket)
			: m_cancelTokenSource(), m_dataReaderTaskGroup(m_cancelTokenSource.get_token())
		{
			m_isListening = false;

			m_dataReader = ref new DataReader(socket->InputStream);
			m_dataReader->ByteOrder = ByteOrder::LittleEndian;
			m_dataReader->UnicodeEncoding = UnicodeEncoding::Utf8;

			m_dataWriter = ref new DataWriter(socket->OutputStream);
			m_dataWriter->ByteOrder = ByteOrder::LittleEndian;
			m_dataWriter->UnicodeEncoding = UnicodeEncoding::Utf8;

			m_socket = socket;

			m_dataReceivedCallback = EmptyDataReceivedCallback;
			m_clientDisconnectedCallback = EmptySocketDisconnectedCallback;
		}

		//
		// Sends data to the connected client.
		//
		GAMESTEST_ANNOTATE_HRESULT_METHOD HRESULT SendData(const std::wstring& data)
		{
			// It is possible that more than one server could be trying to send a response
			// at the same time.
			concurrency::critical_section::scoped_lock lock(m_sendDataLock);

			HRESULT hr = S_OK;
			try
			{
				Platform::StringReference dataString(data.c_str(), data.length());
				Write7BitEncodedInt(m_dataWriter->MeasureString(dataString));
				m_dataWriter->WriteString(dataString);
				create_task(m_dataWriter->StoreAsync()).wait();
			}
			catch(Platform::COMException^ ex)
			{
				hr = ex->HResult;
				return hr;
			}

			return hr;
		}

		//
		// Start listening for incoming data and connection events.
		//	dataReceivedCallback: Function to invoke when data is received.
		//	clientDisconnectedCallback">Function to invoke when a client disconnects
		//
		HRESULT StartListening(DataReceivedCallback dataReceivedCallback, SocketDisconnectedCallback clientDisconnectedCallback)
		{
			m_dataReceivedCallback = dataReceivedCallback;
			m_clientDisconnectedCallback = clientDisconnectedCallback;

			if(!m_isListening)
			{
				AcceptIncomingRequestsAsync();
				m_isListening = true;
			}

			return S_OK;
		}

	private:
		task_group m_dataReaderTaskGroup;

		//
		// This exception is raised when we fail to read the expected number of bytes
		// from the incoming network stream.
		//
		class ReadFailureException : public std::exception
		{
		public:
			ReadFailureException(const char* const& message)
				: std::exception(message)
			{
			}
		};

		WinRtStringSocketT(const WinRtStringSocketT&) { }
		WinRtStringSocketT& operator=(WinRtStringSocketT&) { }

		// This function came straight from the .NET framework's BinaryWriter class.
		// https://codesearch/FileView.aspx?filename=%5c%5ctk5codeindex%5cdotnet%5c8.0%5cdevdiv%5cdepot%5cdevdiv%5creleases%5cwhidbey%5credbits%5cndp%5cclr%5csrc%5cbcl%5csystem%5cio%5cbinarywriter.cs%5c1%5cbinarywriter.cs&term="binarywriter"
		void Write7BitEncodedInt(int value) 
		{
			// Write out an int 7 bits at a time.  The high bit of the byte, 
			// when on, tells reader to continue reading more bytes. 
			UINT32 v = (UINT32) value;   // support negative numbers
			while (v >= 0x80) 
			{ 
				m_dataWriter->WriteByte((BYTE)(v | 0x80));
				v >>= 7;
			}
			m_dataWriter->WriteByte((BYTE)v);
		}

		// This function came pretty much straight from the .NET framework's BinaryReader class.
		// https://codesearch/FileView.aspx?filename=%5c%5ctk5codeindex%5cdotnet%5c8.0%5cdevdiv%5cdepot%5cdevdiv%5creleases%5cwhidbey%5credbits%5cndp%5cclr%5csrc%5cbcl%5csystem%5cio%5cbinaryreader.cs%5c1%5cbinaryreader.cs&term="binaryreader"
		int Read7BitEncodedInt(concurrency::cancellation_token cancelToken) 
		{
			// Read out an Int32 7 bits at a time.  The high bit 
			// of the byte when on means to continue reading more bytes.
			int count = 0;
			int shift = 0;
			BYTE b = 0;
			do 
			{
				// Check for a corrupted stream.  Read a max of 5 bytes. 
				// In a future version, add a DataFormatException. 
				if (shift == 5 * 7)  // 5 bytes max per Int32, shift += 7
					return -1;

				b = ReadByte(cancelToken);
				count |= (b & 0x7F) << shift; 
				shift += 7;
			} while ((b & 0x80) != 0); 
			return count; 
		}

		BYTE ReadByte(concurrency::cancellation_token cancelToken)
		{
			const int NumBytesToLoad = 1;
			auto byteLoader = create_task(m_dataReader->LoadAsync(NumBytesToLoad), cancelToken);
			unsigned __int32 numBytesLoaded = byteLoader.get();
			if(numBytesLoaded != NumBytesToLoad)
			{
				throw ReadFailureException("Failed to read a byte from the network stream.");
			}
			else
			{
				return m_dataReader->ReadByte();
			}
		}

		void AcceptIncomingRequestsAsync()
		{
			auto cancellationToken = m_cancelTokenSource.get_token();

			m_dataReaderTaskGroup.run([=]()
			{
				while(!concurrency::is_current_task_group_canceling())
				{
					unsigned __int32 bufferLength = 0;
					try
					{
						bufferLength = Read7BitEncodedInt(cancellationToken);
					}
					catch (Platform::COMException^)
					{
						RaiseClientDisconnectedEvent(m_clientDisconnectedCallback, this);
						return;
					}
					catch (ReadFailureException&)
					{
						RaiseClientDisconnectedEvent(m_clientDisconnectedCallback, this);
						return;
					}
					catch (const task_canceled&)
					{
						return;
					}

					if (bufferLength == -1)
					{
						RaiseClientDisconnectedEvent(m_clientDisconnectedCallback, this);
						return;
					}

					auto stringLoader = create_task(m_dataReader->LoadAsync(bufferLength), cancellationToken);

					unsigned __int32 numBytesLoadedByStringLoader = 0;
					try
					{
						numBytesLoadedByStringLoader = stringLoader.get();
					}
					catch (Platform::COMException^)
					{
						RaiseClientDisconnectedEvent(m_clientDisconnectedCallback, this);
						return;
					}
					catch (const task_canceled&)
					{
						return;
					}

					if (numBytesLoadedByStringLoader != bufferLength)
					{
						RaiseClientDisconnectedEvent(m_clientDisconnectedCallback, this);
						return;
					}

					Platform::String^ incomingString = m_dataReader->ReadString(bufferLength);

					if (incomingString == L"RpcSocketClientDisconnect")
					{
						RaiseClientDisconnectedEvent(m_clientDisconnectedCallback, this);
						return;
					}
					else
					{
						RaiseDataReceivedEvent(m_dataReceivedCallback, this, std::wstring(incomingString->Data(), incomingString->Length()));
					}
				}
			});
		}

		//
		// Stop listening for incoming data and connection events.
		//
		void Shutdown()
		{
			if(m_isListening)
			{
				m_dataReceivedCallback = EmptyDataReceivedCallback;
				m_clientDisconnectedCallback = EmptySocketDisconnectedCallback;
				
				m_cancelTokenSource.cancel();
				m_dataReaderTaskGroup.wait();
				
				m_isListening = false;
			}
		}

		static void RaiseDataReceivedEvent(DataReceivedCallback& callback, IRpcStringSocket* client, std::wstring data)
		{
			create_task([=]()
			{
				callback(client, data); 
			});
		}

		static void RaiseClientDisconnectedEvent(SocketDisconnectedCallback& callback, IRpcStringSocket* client)
		{
			create_task([=]()
			{
				callback(client);
			});
		}
	};

	typedef WinRtStringSocketT<Empty> WinRtStringSocket;

}}}}}

#endif // _DURANGO